package com.kingsoft.dc.khaos.module.spark.preprocess.transform

import com.kingsoft.dc.khaos.KhaosContext
import com.kingsoft.dc.khaos.innertype.Schema
import com.kingsoft.dc.khaos.metadata.{Dependency, KhaosStructField}
import com.kingsoft.dc.khaos.module.spark.metadata.preprocess.transform.SampleInfo
import com.kingsoft.dc.khaos.module.spark.util.DataframeUtils
import org.apache.spark.sql.DataFrame
import org.json4s.DefaultFormats
import org.json4s.jackson.JsonMethods.parse

/**
  * 采样算子
  * create by goosoog 2019/7/25.
  *
  */
class Sample extends TransformStrategy {

  private var sampleInfo: SampleInfo = _

  override def exec(kc: KhaosContext,
                    module_id: String,
                    config: String,
                    dependences: Seq[Dependency],
                    targets: Seq[Dependency]): Seq[(String, DataFrame)] = {
    //获取配置信息
    implicit val formats = DefaultFormats
    sampleInfo = parse(config, true).extract[SampleInfo]
    val (schema, dataFrame) = kc.structData[DataFrame](dependences.head.getEdge())

    var percent = sampleInfo.sample_rate.toDouble
    if (percent > 0 && percent <= 100) {
      percent = percent / 100.00
    } else {
      throw new Exception(s"采样比例不能为负数或者大于100 sample_rate=${percent}%")
    }
    var finalDF = dataFrame.sample(percent)

    if ("true".equals(sampleInfo.advanced_options.cache_option.on_off)) {
      finalDF = DataframeUtils.cacheDataframe(finalDF, this)
    }
    if ("true".equals(sampleInfo.advanced_options.partition.on_off)) {
      finalDF = DataframeUtils.repartionDataframe(finalDF, sampleInfo.advanced_options.partition.nums.toInt, this)
    }
    addResult(targets.head, finalDF)
  }

  //21202F2938212B3E22272626252E434D
  override def schema(kc: KhaosContext,
                      module_id: String,
                      config: String,
                      dependencies: Seq[Dependency]): Schema = {

    val schema = if (sampleInfo.extract_fields.isEmpty) {
      kc.schemaChannel.getSchema(dependencies.head.getSource())
    } else {
      val schemaList = sampleInfo.extract_fields.map { optInfo =>
        KhaosStructField(optInfo.field, optInfo.data_type)
      }
      new Schema(schemaList)
    }
    schema
  }
}
